from camel.models import ModelFactory
from camel.societies import RolePlaying
from camel.types import ModelPlatformType
from dotenv import load_dotenv
import os

load_dotenv()

Model_Type = os.getenv("MODEL_TYPE")
Model_Api = os.getenv("ZHIPU_API_KEY")
Model_Url = os.getenv("MODEL_URL")
model = ModelFactory.create(
    model_platform=ModelPlatformType.OPENAI,
    model_type=Model_Type,
    api_key=Model_Api,
    url=Model_Url
)
# 制定任务
task_kwargs = {
    'task_prompt': "完成一个详尽的合理的针对于洛阳的三日游旅行规划,全程使用中文沟通",
    'with_task_specify': True,
    'task_specify_agent_kwargs': {
        'model': model,
    }
}
# 创建AI用户

user_role_kwargs = {
    'user_role_name': '平常很少旅游的小白，但是想去洛阳旅游，全程使用中文沟通',
    'user_agent_kwargs': {
        'model': model,
    }
}

assistant_role_kwargs = {
    'assistant_role_name': '经验丰富的洛阳导游，能够针对不同性格的用户给出合理的旅行建议，全程使用中文沟通',
    'assistant_agent_kwargs': {
        'model': model,
    }
}
# 构建society智能体
society = RolePlaying(
    **task_kwargs,
    **user_role_kwargs,
    **assistant_role_kwargs
)


# 检查任务是否终止
def is_terminated(response):
    if response.terminated:
        role = response.msg.role_type.name
        reason = response.info['termination_reasons']
        print(f'AI{role} terminated due to {reason}')
    return response.terminated


# 运行
def run(society, round_limit: int):
    # 创建初始对话
    input_msg = society.init_chat()
    # 开始对话
    for i in range(round_limit):
        # 通过初始对话，开始第一轮对话
        assistant_response, user_response = society.step(input_msg)

        # 检查任务是否终止
        if is_terminated(user_response) or is_terminated(assistant_response):
            break

        # 输出结果
        print(f'[AI User] {user_response.msg.content}.\n')
        # 再次通过CAMEL_TASK_DONE字段判断任务是否结束
        if 'CAMEL_TASK_DONE' in user_response.msg.content:
            break
        print(f"[AI Assistant] {assistant_response.msg.content}.\n")

        # 把本轮消息赋给input_msg以便下轮对话,User/Assistant均可
        input_msg = assistant_response.msg.content
    return None


if __name__ == '__main__':
    run(society, 3)
